import torch
import torch.nn as nn
import torch.nn.functional as F
from matplotlib import pyplot as plt

from src.SDM_module import SDM
from src.transformer import Block1
from src.dsc import IDSC


class decoder(nn.Module):
    def __init__(self, num_class, dims):
        super().__init__()
        channels = [16, 32, 64, 128, 256, 512]

        self.p = nn.PixelShuffle(2)


        self.c1 = IDSC(dims[3]//4 + dims[2], dims[2])
        self.b1 = Block1(dims[2], 8, 4, 0.2)


        self.c2 = IDSC(dims[2]//4 + dims[1], dims[1])
        self.b2 = Block1(dims[1], 8, 4, 0.3)


        self.c3 = IDSC(dims[1]//4 + dims[0], dims[0])
        self.b3 = Block1(dims[0], 8, 8, 0.4)


        self.head = nn.Sequential(nn.PixelShuffle(4),
                                  nn.Conv2d(dims[0]//16, num_class, 1))

        self.sdn1 = SDM(channels[4], channels[5])
        self.sdn2 = SDM(channels[3], channels[4])
        self.sdn3 = SDM(channels[2], channels[3])

        self.rc1 = nn.Conv2d(768, 384, kernel_size=1, stride=1, bias=False)
        self.rc2 = nn.Conv2d(384, 192, kernel_size=1, stride=1, bias=False)
        self.rc3 = nn.Conv2d(192, 96, kernel_size=1, stride=1, bias=False)

    def forward(self, x1, x2, x3, x4):
        # print(x1.shape, x2.shape, x3.shape, x4.shape)
        # torch.Size([24, 64, 64, 64]) torch.Size([24, 128, 32, 32])
        # torch.Size([24, 256, 16, 16]) torch.Size([24, 512, 8, 8])

        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
        # print("40", x4.shape) torch.Size([24, 512, 16, 16])
        # x4 = self.p(x4)
        # print(x4.shape) torch.Size([24, 128, 16, 16])
        x3 = self.sdn1(x3, x4)
        # print("45", x3.shape) torch.Size([24, 256, 16, 16])
        x3 = torch.cat([x3, x4], dim=1)
        x3 = self.rc1(x3)
        x3 = self.c1(x3)
        b, c, h, w = x3.shape
        x3 = self.b1(x3.reshape(b, c, -1).permute(0, 2, 1))
        x3 = x3.permute(0, 2, 1).reshape(b, c, h, w)
        # print("55", x3.shape) torch.Size([24, 256, 16, 16])

        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
        # print(x2.shape, x3.shape) torch.Size([24, 128, 32, 32]) torch.Size([24, 256, 32, 32])
        # x3 = self.p(x3)
        x2 = self.sdn2(x2, x3)
        # print("62", x2.shape) torch.Size([24, 128, 32, 32])
        x2 = torch.cat([x2, x3], dim=1)
        x2 = self.rc2(x2)
        x2 = self.c2(x2)
        b, c, h, w = x2.shape
        x2 = self.b2(x2.reshape(b, c, -1).permute(0, 2, 1))
        x2 = x2.permute(0, 2, 1).reshape(b, c, h, w)

        x2 = F.interpolate(x2, scale_factor=2, mode='bilinear', align_corners=False)
        # x2 = self.p(x2)
        x1 = self.sdn3(x1, x2)
        x1 = torch.cat([x1, x2], dim=1)
        x1 = self.rc3(x1)
        x1 = self.c3(x1)
        b, c, h, w = x1.shape
        x1 = self.b3(x1.reshape(b, c, -1).permute(0, 2, 1))
        x1 = x1.permute(0, 2, 1).reshape(b, c, h, w)

        out = self.head(x1)
        return out